#!/usr/bin/env python3

"""
uses TensorRT's Python api to make inferences.
"""

import ctypes
import os
import shutil
import random
import sys
import threading
import time
import cv2
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt

from sort3d import Sort3D

import rospy
from sensor_msgs.msg import Image
from geometry_msgs.msg import Vector3Stamped
from geometry_msgs.msg import PoseStamped
from geometry_msgs.msg import Point
from traj_utils.msg import DetectPose
from yolo_detect.msg import YoloResult
from cv_bridge import CvBridge, CvBridgeError
import tf.transformations as transformations

CONF_THRESH = 0.5
IOU_THRESHOLD = 0.4
LEN_ALL_RESULT = 38001
LEN_ONE_RESULT = 38

SORT_MAX_AGE = 90
SORT_MIN_HITS = 0
SORT_DIST_THRESH = 2.0

IS_DISPLAY = False
IS_RECORD = True
IS_SAVE_IMAGE = False

categories = [
    "person",
    "bicycle",
    "car",
    "motorcycle",
    "airplane",
    "bus",
    "train",
    "truck",
    "boat",
    "traffic light",
    "fire hydrant",
    "stop sign",
    "parking meter",
    "bench",
    "bird",
    "cat",
    "dog",
    "horse",
    "sheep",
    "cow",
    "elephant",
    "bear",
    "zebra",
    "giraffe",
    "backpack",
    "umbrella",
    "handbag",
    "tie",
    "suitcase",
    "frisbee",
    "skis",
    "snowboard",
    "sports ball",
    "kite",
    "baseball bat",
    "baseball glove",
    "skateboard",
    "surfboard",
    "tennis racket",
    "bottle",
    "wine glass",
    "cup",
    "fork",
    "knife",
    "spoon",
    "bowl",
    "banana",
    "apple",
    "sandwich",
    "orange",
    "broccoli",
    "carrot",
    "hot dog",
    "pizza",
    "donut",
    "cake",
    "chair",
    "couch",
    "potted plant",
    "bed",
    "dining table",
    "toilet",
    "tv",
    "laptop",
    "mouse",
    "remote",
    "keyboard",
    "cell phone",
    "microwave",
    "oven",
    "toaster",
    "sink",
    "refrigerator",
    "book",
    "clock",
    "vase",
    "scissors",
    "teddy bear",
    "hair drier",
    "toothbrush",
]


def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    """
    description: Plots one bounding box on image img,
                 this function comes from YoLov5 project.
    param:
        x:      a box likes [x1,y1,x2,y2]
        img:    a opencv image object
        color:  color to draw rectangle, such as (0,255,0)
        label:  str
        line_thickness: int
    return:
        no return

    """
    tl = (
        line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    )  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
        cv2.putText(
            img,
            label,
            (c1[0], c1[1] - 2),
            0,
            tl / 3,
            [225, 255, 255],
            thickness=tf,
            lineType=cv2.LINE_AA,
        )


class YoLov5TRT(object):
    """
    description: A YOLOv5 class that warps TensorRT ops, preprocess and postprocess ops.
    """

    def __init__(self, engine_file_path):
        # Initialize SORT
        self.tracker = Sort3D(
            max_age=SORT_MAX_AGE,
            min_hits=SORT_MIN_HITS,
            dist_thresh=SORT_DIST_THRESH,
        )

        # parameters
        self.is_warmingup = True

        # Create a Context on this device,
        self.ctx = cuda.Device(0).make_context()
        stream = cuda.Stream()
        TRT_LOGGER = trt.Logger(trt.Logger.INFO)
        runtime = trt.Runtime(TRT_LOGGER)

        # Deserialize the engine from file
        with open(engine_file_path, "rb") as f:
            engine = runtime.deserialize_cuda_engine(f.read())
        context = engine.create_execution_context()

        host_inputs = []
        cuda_inputs = []
        host_outputs = []
        cuda_outputs = []
        bindings = []

        for binding in engine:
            tensor_shape = engine.get_tensor_shape(binding)
            print("bingding:", binding, tensor_shape)
            size = trt.volume(tensor_shape) * engine.max_batch_size
            dtype = trt.nptype(engine.get_tensor_dtype(binding))
            host_mem = cuda.pagelocked_empty(size, dtype)
            cuda_mem = cuda.mem_alloc(host_mem.nbytes)
            bindings.append(int(cuda_mem))
            if engine.binding_is_input(binding):
                self.input_w = tensor_shape[-1]
                self.input_h = tensor_shape[-2]
                host_inputs.append(host_mem)
                cuda_inputs.append(cuda_mem)
            else:
                host_outputs.append(host_mem)
                cuda_outputs.append(cuda_mem)

        # Store
        self.stream = stream
        self.context = context
        self.engine = engine
        self.host_inputs = host_inputs
        self.cuda_inputs = cuda_inputs
        self.host_outputs = host_outputs
        self.cuda_outputs = cuda_outputs
        self.bindings = bindings
        self.batch_size = engine.max_batch_size

        # init ros
        self.bridge = CvBridge()
        self.depth_image = None
        self.angular_velocity = 0.0
        self.have_depth = False
        self.pose = None
        self.color_time = None
        self.image_count = 0

        if IS_SAVE_IMAGE:
            self.path = os.path.dirname(__file__)
            self.path = os.path.join(self.path, "../images")
            self.path_yolo = os.path.join(self.path, "../output")
            os.makedirs(self.path, exist_ok=True)
            os.makedirs(self.path_yolo, exist_ok=True)

        rospy.init_node("yolov5_trt", anonymous=True)
        rospy.Subscriber("/camera/color/image_raw", Image, self.color_cb)
        rospy.Subscriber(
            "/camera/aligned_depth_to_color/image_raw", Image, self.depth_cb
        )
        rospy.Subscriber("/servo/state", Vector3Stamped, self.servo_cb)
        rospy.Subscriber("/camera/pose", PoseStamped, self.pose_cb)

        self.pub_pos = rospy.Publisher("/detect/objects", DetectPose, queue_size=10)
        self.pub_img = rospy.Publisher("/detect/image", Image, queue_size=1)

        # camera intrinsic parameters
        self.fx = 610.8905029296875
        self.fy = 610.8898315429688
        self.cx = 322.75732421875
        self.cy = 250.6327667236328
        self.depth_calib = np.array([0.00567627, 1.02187439, -0.04152119, 0.00138622])
        self.depth_calib = self.depth_calib[::-1]

    def color_cb(self, data):
        if not self.have_depth:
            return
        try:
            cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
        except CvBridgeError as e:
            print(e)
        else:
            self.color_time = data.header.stamp

            if IS_SAVE_IMAGE:
                image_path = os.path.join(self.path, f"red_{self.image_count:04d}.jpg")
                cv2.imwrite(image_path, cv_image)
                rospy.loginfo(f"Saved image: {image_path}")

            output_frame, use_time = self.infer(cv_image)
            print("----------------time->{:.2f}ms".format(use_time * 1000))
            self.have_depth = False

            if IS_SAVE_IMAGE:
                image_path = os.path.join(
                    self.path_yolo, f"red_{self.image_count:04d}.jpg"
                )
                cv2.imwrite(image_path, output_frame)
                self.image_count += 1

            # print("time->{:.2f}ms".format(use_time * 1000))
            if IS_DISPLAY:
                cv2.imshow("YOLOv5 Detection", output_frame)
                cv2.waitKey(1)  # Small delay to update the display

    def depth_cb(self, data):
        # only update depth image when angular velocity is 0
        if self.angular_velocity != 0.0 or self.is_warmingup:
            return
        try:
            self.depth_image = self.bridge.imgmsg_to_cv2(data, "32FC1")
            self.have_depth = True
        except CvBridgeError as e:
            print(e)

    def servo_cb(self, data):
        self.angular_velocity = data.vector.y

    def pose_cb(self, data):
        self.pose = data

    def convert_to_world_frame(self, X, Y, Z):
        if self.pose is None:
            return np.array([0.0, 0.0, 0.0])

        position = self.pose.pose.position
        orientation = self.pose.pose.orientation

        quaternion = [orientation.x, orientation.y, orientation.z, orientation.w]
        rotation_matrix = transformations.quaternion_matrix(quaternion)

        camera_to_object_vector = np.array([X, Y, Z, 1])  # 齐次坐标

        world_position = rotation_matrix.dot(camera_to_object_vector)
        world_position = world_position[:3] + np.array(
            [position.x, position.y, position.z]
        )

        return world_position

    def convert_to_camera_frame(self, X, Y, Z):
        if self.pose is None:
            return np.array([0.0, 0.0, 0.0])

        p = self.pose.pose.position
        q = self.pose.pose.orientation
        p_cam_world = np.array([p.x, p.y, p.z])
        p_obj_world = np.array([X, Y, Z])
        q_world2cam = [q.x, q.y, q.z, q.w]
        m_world2cam = transformations.quaternion_matrix(q_world2cam)
        m_cam2world = m_world2cam.transpose()
        p_world = p_obj_world - p_cam_world
        p_world = np.append(p_world, 1)
        return m_cam2world.dot(p_world)[:3]

    def infer(self, image_raw):
        pre_time = time.time()
        # Make self the active context, pushing it on top of the context stack.
        self.ctx.push()
        # Restore
        stream = self.stream
        context = self.context
        engine = self.engine
        host_inputs = self.host_inputs
        cuda_inputs = self.cuda_inputs
        host_outputs = self.host_outputs
        cuda_outputs = self.cuda_outputs
        bindings = self.bindings

        # Image preprocessing
        input_image, image_raw, origin_h, origin_w = self.preprocess_image(image_raw)
        input_image = np.ascontiguousarray(input_image)

        # Copy input image to host buffer
        np.copyto(host_inputs[0], input_image.ravel())
        start = time.time()
        # Transfer input data  to the GPU.
        cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
        # Run inference.
        context.execute_async(
            batch_size=1, bindings=bindings, stream_handle=stream.handle
        )
        # Transfer predictions back from the GPU.
        cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
        # Synchronize the stream
        stream.synchronize()
        end_infer = time.time()
        # Remove any context from the top of the context stack, deactivating it.
        self.ctx.pop()
        # Here we use the first row of output in that batch_size = 1
        # Output handling
        output = host_outputs[0]
        result_boxes, result_scores, result_classid = self.post_process(
            output, origin_h, origin_w
        )

        if self.is_warmingup:
            return image_raw, end_infer - start

        end_post = time.time()

        # # Draw rectangles and labels on the original image
        # if IS_DISPLAY or IS_RECORD or IS_SAVE_IMAGE:
        #     for j in range(len(result_boxes)):
        #         box = result_boxes[j]
        #         plot_one_box(
        #             box,
        #             image_raw,
        #             # label="{}:{:.2f}".format("uav", result_scores[j]),
        #             label="{}:{:.2f}".format(
        #                 categories[int(result_classid[j])], result_scores[j]
        #             ),
        #         )

        # xyzs = np.zeros((len(result_boxes), 3))
        # for i in range(len(result_boxes)):
        #     xyzs[i] = self.xyxy2xyz(result_boxes[i])

        # # Update the tracker with the new frame detections
        # # tracks = self.tracker.update(xyzs)  # TODO: comment temporarily
        # zero_column = np.zeros((len(result_boxes), 1))
        # tracks = np.hstack((xyzs, zero_column))

        # # Publish the detection results
        # msg = DetectPose()
        # msg.header.stamp = self.color_time

        # for track in tracks:
        #     x, y, z, track_id = track
        #     track_id = int(track_id)
        #     pose = Point()
        #     msg.ids.append(track_id)
        #     pose.x = x
        #     pose.y = y
        #     pose.z = z
        #     msg.poses.append(pose)

        #     if IS_DISPLAY or IS_RECORD:
        #         center_x, center_y = self.xyz2xy(x, y, z)
        #         label = f"{track_id}"
        #         label += f" X: {x:.2f}, Y: {y:.2f}, Z: {z:.2f}"
        #         cv2.putText(
        #             image_raw,
        #             label,
        #             (center_x, center_y),
        #             cv2.FONT_HERSHEY_SIMPLEX,
        #             0.5,
        #             (0, 0, 255),
        #             2,
        #             cv2.LINE_AA,
        #         )

        # if not msg.ids == []:
        #     self.pub_pos.publish(msg)

        # if IS_RECORD and self.pub_img.get_num_connections() > 0:
        #     self.pub_img.publish(self.bridge.cv2_to_imgmsg(image_raw, "bgr8"))

        end_all = time.time()

        if True or end_all - start > 0.1:
            print("pre time  ->{:.2f}ms".format((start - pre_time) * 1000))
            print("infer time->{:.2f}ms".format((end_infer - start) * 1000))
            print("post time ->{:.2f}ms".format((end_post - end_infer) * 1000))
            print("left time ->{:.2f}ms".format((end_all - end_post) * 1000))

        return image_raw, end_all - pre_time

    def destroy(self):
        # Remove any context from the top of the context stack, deactivating it.
        self.ctx.pop()

    def get_raw_image_zeros(self, image_path_batch=None):
        """
        description: Ready data for warmup
        """
        return np.zeros([self.input_h, self.input_w, 3], dtype=np.uint8)

    def preprocess_image(self, raw_bgr_image):
        """
        description: Convert BGR image to RGB,
                     resize and pad it to target size, normalize to [0,1],
                     transform to NCHW format.
        param:
            input_image_path: str, image path
        return:
            image:  the processed image
            image_raw: the original image
            h: original height
            w: original width
        """
        image_raw = raw_bgr_image
        h, w, c = image_raw.shape
        image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB)
        # Calculate widht and height and paddings
        r_w = self.input_w / w
        r_h = self.input_h / h
        if r_h > r_w:
            tw = self.input_w
            th = int(r_w * h)
            tx1 = tx2 = 0
            ty1 = int((self.input_h - th) / 2)
            ty2 = self.input_h - th - ty1
        else:
            tw = int(r_h * w)
            th = self.input_h
            tx1 = int((self.input_w - tw) / 2)
            tx2 = self.input_w - tw - tx1
            ty1 = ty2 = 0
        # Resize the image with long side while maintaining ratio
        image = cv2.resize(image, (tw, th))
        # Pad the short side with (128,128,128)
        image = cv2.copyMakeBorder(
            image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, None, (128, 128, 128)
        )
        image = image.astype(np.float32)
        # Normalize to [0,1]
        image /= 255.0
        # HWC to CHW format:
        image = np.transpose(image, [2, 0, 1])
        # CHW to NCHW format
        image = np.expand_dims(image, axis=0)
        # Convert the image to row-major order, also known as "C order":
        image = np.ascontiguousarray(image)
        return image, image_raw, h, w

    def xywh2xyxy(self, origin_h, origin_w, x):
        """
        description:    Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
        param:
            origin_h:   height of original image
            origin_w:   width of original image
            x:          A boxes numpy, each row is a box [center_x, center_y, w, h]
        return:
            y:          A boxes numpy, each row is a box [x1, y1, x2, y2]
        """
        y = np.zeros_like(x)
        r_w = self.input_w / origin_w
        r_h = self.input_h / origin_h
        if r_h > r_w:
            y[:, 0] = x[:, 0] - x[:, 2] / 2
            y[:, 2] = x[:, 0] + x[:, 2] / 2
            y[:, 1] = x[:, 1] - x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
            y[:, 3] = x[:, 1] + x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
            y /= r_w
        else:
            y[:, 0] = x[:, 0] - x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
            y[:, 2] = x[:, 0] + x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
            y[:, 1] = x[:, 1] - x[:, 3] / 2
            y[:, 3] = x[:, 1] + x[:, 3] / 2
            y /= r_h

        return y

    def post_process(self, output, origin_h, origin_w):
        """
        description: postprocess the prediction
        param:
            output:     A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...]
            origin_h:   height of original image
            origin_w:   width of original image
        return:
            result_boxes: finally boxes, a boxes numpy, each row is a box [x1, y1, x2, y2]
            result_scores: finally scores, a numpy, each element is the score correspoing to box
            result_classid: finally classid, a numpy, each element is the classid correspoing to box
        """
        # Get the num of boxes detected
        num = int(output[0])
        # Reshape to a two dimentional ndarray
        pred = np.reshape(output[1:], (-1, LEN_ONE_RESULT))[:num, :]
        pred = pred[:, :6]
        # Do nms
        boxes = self.non_max_suppression(
            pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD
        )
        result_boxes = boxes[:, :4] if len(boxes) else np.array([])
        result_scores = boxes[:, 4] if len(boxes) else np.array([])
        result_classid = boxes[:, 5] if len(boxes) else np.array([])
        return result_boxes, result_scores, result_classid

    def bbox_iou(self, box1, box2, x1y1x2y2=True):
        """
        description: compute the IoU of two bounding boxes
        param:
            box1: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h))
            box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h))
            x1y1x2y2: select the coordinate format
        return:
            iou: computed iou
        """
        if not x1y1x2y2:
            # Transform from center and width to exact coordinates
            b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
            b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
            b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
            b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
        else:
            # Get the coordinates of bounding boxes
            b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
            b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]

        # Get the coordinates of the intersection rectangle
        inter_rect_x1 = np.maximum(b1_x1, b2_x1)
        inter_rect_y1 = np.maximum(b1_y1, b2_y1)
        inter_rect_x2 = np.minimum(b1_x2, b2_x2)
        inter_rect_y2 = np.minimum(b1_y2, b2_y2)
        # Intersection area
        inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) * np.clip(
            inter_rect_y2 - inter_rect_y1 + 1, 0, None
        )
        # Union Area
        b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
        b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)

        iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

        return iou

    def non_max_suppression(
        self, prediction, origin_h, origin_w, conf_thres=0.5, nms_thres=0.4
    ):
        """
        description: Removes detections with lower object confidence score than 'conf_thres' and performs
        Non-Maximum Suppression to further filter detections.
        param:
            prediction: detections, (x1, y1, x2, y2, conf, cls_id)
            origin_h: original image height
            origin_w: original image width
            conf_thres: a confidence threshold to filter detections
            nms_thres: a iou threshold to filter detections
        return:
            boxes: output after nms with the shape (x1, y1, x2, y2, conf, cls_id)
        """
        # Get the boxes that score > CONF_THRESH
        boxes = prediction[prediction[:, 4] >= conf_thres]
        # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2]
        boxes[:, :4] = self.xywh2xyxy(origin_h, origin_w, boxes[:, :4])
        # clip the coordinates
        boxes[:, 0] = np.clip(boxes[:, 0], 0, origin_w - 1)
        boxes[:, 2] = np.clip(boxes[:, 2], 0, origin_w - 1)
        boxes[:, 1] = np.clip(boxes[:, 1], 0, origin_h - 1)
        boxes[:, 3] = np.clip(boxes[:, 3], 0, origin_h - 1)
        # Object confidence
        confs = boxes[:, 4]
        # Sort by the confs
        boxes = boxes[np.argsort(-confs)]
        # Perform non-maximum suppression
        keep_boxes = []
        while boxes.shape[0]:
            large_overlap = (
                self.bbox_iou(np.expand_dims(boxes[0, :4], 0), boxes[:, :4]) > nms_thres
            )
            label_match = boxes[0, -1] == boxes[:, -1]
            # Indices of boxes with lower confidence scores, large IOUs and matching labels
            invalid = large_overlap & label_match
            keep_boxes += [boxes[0]]
            boxes = boxes[~invalid]
        boxes = np.stack(keep_boxes, 0) if len(keep_boxes) else np.array([])
        return boxes

    def get_depth(self, bbox, threshold=0.2, n=3):
        """
        param:
        bbox: 矩形框，格式为 (x1, y1, x2, y2)
        threshold: 深度阈值，过滤小于此值的深度
        n: 选择最小深度的数量

        返回:
        average_depth: n个最小深度的平均值
        """
        x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
        center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
        small_width, small_height = 8, 6
        small_width = abs(x2 - x1) if abs(x2 - x1) < small_width else small_width
        small_height = abs(y2 - y1) if abs(y2 - y1) < small_height else small_height

        # 计算小矩形框的边界
        small_x1 = int(max(0, center_x - small_width // 2))
        small_x2 = int(min(self.depth_image.shape[1], center_x + small_width // 2))
        small_y1 = int(max(0, center_y - small_height // 2))
        small_y2 = int(min(self.depth_image.shape[0], center_y + small_height // 2))

        # 提取小矩形框内的深度值
        small_depth_area = self.depth_image[small_y1:small_y2, small_x1:small_x2]

        # 过滤深度值
        filtered_depths = small_depth_area[
            (small_depth_area >= threshold) & (np.isfinite(small_depth_area))
        ]

        # 检查是否有足够的深度值
        if filtered_depths.size < n:
            return 16.0  # TODO: 如果不够n个，则返回max_depth
            # small_depth_area = self.depth_image[y1:y2, x1:x2]
            # filtered_depths = small_depth_area[
            #     (small_depth_area >= threshold) & (np.isfinite(small_depth_area))
            # ]
            # if filtered_depths.size < n:
            #     print(
            #         "[YOLO DETECT]: Not enough depth values, size: ",
            #         filtered_depths.size,
            #     )
            #     return 16.0  # TODO: 如果不够n个，则返回max_depth

        # 选取最小的n个深度
        smallest_depths = np.partition(filtered_depths, n - 1)[:n]

        # 计算平均值
        average_depth = np.mean(smallest_depths)

        # Depth calibration ,convert to meters
        # average_depth = np.polyval(self.depth_calib, average_depth * 0.001)
        average_depth = average_depth * 0.001

        return average_depth

    def xyxy2xyz(self, xyxy):
        """
        description: Convert bounding box from (x1, y1, x2, y2) to (x, y, z)
        param:
            xyxy: A box coordinate (x1, y1, x2, y2)
        return:
            xyz: A box coordinate (x, y, z)
        """
        x1, y1, x2, y2 = xyxy
        center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
        depth = self.get_depth(xyxy)
        if not np.isnan(depth):
            Z = depth
            X = (center_x - self.cx) * Z / self.fx
            Y = (center_y - self.cy) * Z / self.fy
            return self.convert_to_world_frame(X, Y, Z)
        return np.array([0.0, 0.0, 0.0])

    def xyz2xy(self, x, y, z):
        X, Y, Z = self.convert_to_camera_frame(x, y, z)
        if X == 0.0 and Y == 0.0 and Z == 0.0:
            return 0, 0
        center_x = int(X * self.fx / Z + self.cx)
        center_y = int(Y * self.fy / Z + self.cy)
        return center_x, center_y


class warmUpThread(threading.Thread):
    def __init__(self, yolov5_wrapper):
        threading.Thread.__init__(self)
        self.yolov5_wrapper = yolov5_wrapper

    def run(self):
        image_raw, use_time = self.yolov5_wrapper.infer(
            self.yolov5_wrapper.get_raw_image_zeros()
        )
        print("warm_up->{}, time->{:.2f}ms".format(image_raw.shape, use_time * 1000))


if __name__ == "__main__":
    # load custom plugin and engine
    script_dir = os.path.dirname(__file__)

    # PLUGIN_LIBRARY = os.path.join(script_dir, "../build/libmyplugins.so")
    # engine_file_path = os.path.join(script_dir, "../build/yolov5s.engine")

    # PLUGIN_LIBRARY = os.path.join(script_dir, "../build/libmyplugins_coco.so")
    # engine_file_path = os.path.join(script_dir, "../build/yolov5s_coco.engine")

    PLUGIN_LIBRARY = os.path.join(
        script_dir, "../build/libmyplugins_coco_person_only.so"
    )
    engine_file_path = os.path.join(
        script_dir, "../build/yolov5s_coco_person_only.engine"
    )

    ctypes.CDLL(PLUGIN_LIBRARY)

    if os.path.exists("output/"):
        shutil.rmtree("output/")
    os.makedirs("output/")
    # a YoLov5TRT instance
    yolov5_wrapper = YoLov5TRT(engine_file_path)

    try:
        for i in range(10):
            # create a new thread to do warm_up
            thread1 = warmUpThread(yolov5_wrapper)
            thread1.start()
            thread1.join()

        yolov5_wrapper.is_warmingup = False

        rospy.spin()

    finally:
        cv2.destroyAllWindows()
        yolov5_wrapper.destroy()
